import pandas as pd
from wordcloud import WordCloud, ImageColorGenerator
import matplotlib.pyplot as plt
import gradio as gr
from PIL import Image
import numpy as np
import os

# 读取 txt 文件并转换为 DataFrame
def read_freq_from_txt(file_path):
    words = []
    frequencies = []
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        for line in lines:
            word, freq = line.strip().split(': ')
            words.append(word)
            frequencies.append(int(freq))
        return pd.DataFrame({'词语': words, '词频': frequencies})
    except FileNotFoundError:
        print(f"未找到文件: {file_path}")
        return None
    except ValueError:
        print(f"文件格式有误，请确保每行格式为 '词语: 词频'。")
        return None

path = os.getcwd()
x = os.path.join(path,"data\\mask2.png")
# 生成词云图并保存为图片
def generate_wordcloud(file_path, save_path):
    word_freq = read_freq_from_txt(file_path)
    if word_freq is None:
        return None
    try:
        # 读取一个图片作为词云的形状
        mask_image = np.array(Image.open(x))
        # 定义颜色生成器
        image_colors = ImageColorGenerator(mask_image)

        wordcloud = WordCloud(
            font_path='simhei.ttf',  # 中文字体路径
            background_color='white',
            width=800,
            height=600,
            mask=mask_image,  # 使用图片作为词云的形状
            color_func=image_colors,  # 使用颜色生成器
            max_words=100,  # 最大显示的词语数量
            min_font_size=10,  # 最小字体大小
            max_font_size=100,  # 最大字体大小
            random_state=42  # 随机种子，保证每次生成的词云布局一致
        ).generate_from_frequencies(dict(zip(word_freq['词语'], word_freq['词频'])))

        wordcloud.to_file(save_path)  # 保存图片
        return save_path  # 返回文件路径
    except Exception as e:
        print(f"生成词云图失败: {e}")
        return None

# Gradio 界面(词云图)
with gr.Blocks(title="词云图生成器") as demo:
    gr.Markdown("### 基于 TXT 文件的词云图生成器")

    # 文件输入组件
    txt_file = gr.File(label="上传词频 TXT 文件（每行格式：词语: 词频）")

    # 图片显示组件
    wordcloud_img = gr.Image(label="生成的词云图", type="filepath")

    # 生成按钮
    generate_btn = gr.Button("生成词云图")

    # 按钮点击事件逻辑
    def update_image(file):
        if not file:
            return None
        # 临时保存上传的文件
        temp_path = file.name
        # 生成词云图并保存到指定路径
        save_path = os.path.join(path,"data\\generated_wordcloud.png")
        result = generate_wordcloud(temp_path, save_path)
        return result if result else None

    # 绑定按钮事件
    generate_btn.click(
        fn=update_image,
        inputs=txt_file,
        outputs=wordcloud_img
    )

if __name__ == "__main__":
    demo.launch()